package reservoir_sampling

import (
	"math/rand"
	"time"
)

func GetRandomBox(size int) RandomBox {
	return RandomBox{
		size: size,
		pool: make([]int,size),
		count: 0,
	}
}

type RandomBox struct {
	pool []int
	size int
	count int
}

func (p *RandomBox) Add (n int) {
	p.count++
	if n <= p.size {
		p.pool[p.count-1] = n
	} else {
		if p.Rand(p.count) <= p.size {
			p.pool[p.Rand(p.size)-1] = n
		}
	}
}

func (p *RandomBox) Rand(n int) int {
	// 返回一个 1-n 的数
	time.Sleep(time.Nanosecond)
	rand.Seed(time.Now().UnixNano())
	return rand.Intn(n) + 1
}

func (p *RandomBox) Choices() []int {
	ans := make([]int, p.size)
	for i:=0; i<p.size; i++ {
		ans[i] = p.pool[i]
	}
	return ans
}

